/*
Revised from:
SEAL/eval/cpp at master · wctu/SEAL
https://github.com/wctu/SEAL/tree/master/eval/cpp
*/

#include <Python.h>
#include <iostream>
#include <cstring>  // memset
#include <math.h>   // sqrt and power
using namespace std;

#if PY_MAJOR_VERSION >= 3
long (*PyInt_AsLong)(PyObject *obj) = PyLong_AsLong;
#endif


/*
The input format is H*W numpy array integer labels, where H*W is the image size.
The labels must be from 0, and consecutive integers.
*/


static PyObject* computeASA(PyObject* self, PyObject* args) {
    PyObject* pListSP;      // input
    PyObject* pListGT;      // input
    PyObject* pListError;   // output
    PyObject* pOutputObject;   // output
    PyObject* pItem;
    int returnErrorMap = 0;
    Py_ssize_t nr_pixel_sp;
    Py_ssize_t nr_pixel_gt;

    // Parse Python args
    if(!PyArg_ParseTuple(args, "OO|i", &pListSP, &pListGT, &returnErrorMap)) {
        PyErr_SetString(PyExc_TypeError, "Arg: SP_label_list, GT_label_list, returnErrorMap=0");
        return NULL;
    }

    // Get input image from Python list to a double array
    nr_pixel_sp = PyList_Size(pListSP);
    nr_pixel_gt = PyList_Size(pListGT);
    if(nr_pixel_sp != nr_pixel_gt) {
        PyErr_SetString(PyExc_TypeError, "Number of pixels mismatched");
        printf("nr_pixel_sp = %d\n", int(nr_pixel_sp));
        printf("nr_pixel_gt = %d\n", int(nr_pixel_gt));
        return NULL;
    }

    int nr_pixel = int(nr_pixel_sp);

    int* SP = new int[nr_pixel];
    int* GT = new int[nr_pixel];

    for(int i = 0; i < nr_pixel; ++i) {
        pItem = PyList_GetItem(pListSP, i);
        SP[i] = PyInt_AsLong(pItem);
        pItem = PyList_GetItem(pListGT, i);
        GT[i] = PyInt_AsLong(pItem);
    }

    // Find number of superpixels and groundtruth segments
    int nr_sp = 0;
    int nr_gt = 0;
    for(int i = 0; i < nr_pixel; ++i) {
        if(SP[i] > nr_sp) nr_sp = SP[i];
        if(GT[i] > nr_gt) nr_gt = GT[i];
    }
    ++nr_sp;
    ++nr_gt;

//    cout << "nr_sp: " << nr_sp << endl;
//    cout << "nr_gt: " << nr_gt << endl;

    double asa = 0.0;
    int* sp_hist = new int[nr_sp*nr_gt];    // a 2D table used to record SP and GT overlaps
    int* sp_size = new int[nr_sp];
    int* max_overlap_size = new int[nr_sp];
    memset(sp_hist, 0, sizeof(int)*nr_sp*nr_gt);
    memset(sp_size, 0, sizeof(int)*nr_sp);
    for(int i = 0; i < nr_pixel; ++i) {
        sp_hist[SP[i]*nr_gt + GT[i]] += 1;
        sp_size[SP[i]] += 1;
    }
    for(int i = 0; i < nr_sp; ++i) {
        max_overlap_size[i] = 0;
        for(int j = 0; j < nr_gt; ++j) {
            if(sp_hist[i*nr_gt+j] > max_overlap_size[i]) {
                max_overlap_size[i] = sp_hist[i*nr_gt+j];
            }
        }
//        cout << "max_overlap_size[i] : " << max_overlap_size[i] << endl;
        asa += double(max_overlap_size[i]);
    }
    asa /= double(nr_pixel);

    if(returnErrorMap) {
        int leakage;
        pListError = PyList_New(nr_pixel);
        for(int i = 0; i < nr_pixel; ++i) {
            leakage = sp_size[SP[i]] - max_overlap_size[SP[i]];
            pItem = PyFloat_FromDouble(double(leakage));
            PyList_SetItem(pListError, i, pItem);
        }
        pOutputObject = Py_BuildValue("dO", asa, pListError);
    }
    else {
        pOutputObject = Py_BuildValue("d", asa);
    }

    delete[] SP;
    delete[] GT;
    delete[] sp_hist;
    delete[] sp_size;
    delete[] max_overlap_size;
    // delete[] max_overlap_seg;

    return pOutputObject;
}

static PyObject* computeUE(PyObject* self, PyObject* args) {
    PyObject* pListSP;      // input
    PyObject* pListGT;      // input
    PyObject* pListError;   // output
    PyObject* pOutputObject;   // output
    PyObject* pItem;
    int returnErrorMap = 0;
    Py_ssize_t nr_pixel_sp;
    Py_ssize_t nr_pixel_gt;

    // Parse Python args
    if(!PyArg_ParseTuple(args, "OO|i", &pListSP, &pListGT, &returnErrorMap)) {
        PyErr_SetString(PyExc_TypeError, "Arg: SP_label_list, GT_label_list, returnErrorMap=0");
        return NULL;
    }

    // Get input image from Python list to a double array
    nr_pixel_sp = PyList_Size(pListSP);
    nr_pixel_gt = PyList_Size(pListGT);
    if(nr_pixel_sp != nr_pixel_gt) {
        PyErr_SetString(PyExc_TypeError, "Number of pixels mismatched");
        printf("nr_pixel_sp = %d\n", int(nr_pixel_sp));
        printf("nr_pixel_gt = %d\n", int(nr_pixel_gt));
        return NULL;
    }

    int nr_pixel = int(nr_pixel_sp);

    int* SP = new int[nr_pixel];
    int* GT = new int[nr_pixel];

    for(int i = 0; i < nr_pixel; ++i) {
        pItem = PyList_GetItem(pListSP, i);
        SP[i] = PyInt_AsLong(pItem);
        pItem = PyList_GetItem(pListGT, i);
        GT[i] = PyInt_AsLong(pItem);
    }

    // Find number of superpixels and groundtruth segments: nr_sp, nr_gt
    int nr_sp = 0;
    int nr_gt = 0;
    for(int i = 0; i < nr_pixel; ++i) {
        if(SP[i] > nr_sp) nr_sp = SP[i];
        if(GT[i] > nr_gt) nr_gt = GT[i];
    }
    // include the starting point 0, so ++nr_sp
    ++nr_sp;
    ++nr_gt;

    double ue = 0.0;
    int* gt_hist = new int[nr_gt*nr_sp];    // a 2D table used to record SP and GT overlaps
    int* gt_size = new int[nr_gt];
    int* sp_size = new int[nr_sp];
    int* accu_sp_size = new int[nr_gt];
    memset(gt_hist, 0, sizeof(int)*nr_gt*nr_sp);
    memset(gt_size, 0, sizeof(int)*nr_gt);
    memset(sp_size, 0, sizeof(int)*nr_sp);
    for(int i = 0; i < nr_pixel; ++i) {
        gt_hist[GT[i]*nr_sp + SP[i]] += 1;
        gt_size[GT[i]] += 1;
        sp_size[SP[i]] += 1;
    }
    for(int i = 0; i < nr_gt; ++i) {
        accu_sp_size[i] = 0;
        for(int j = 0; j < nr_sp; ++j) {
            if(gt_hist[i*nr_sp+j] > 0) {
                accu_sp_size[i] += sp_size[j];
            }
        }
//        cout << "i: " << i << " | accu_sp_size: " << accu_sp_size[i] << " | gt_size: " << gt_size[i] << endl;
        ue += double(accu_sp_size[i] - gt_size[i]) / double(gt_size[i]);
    }
    ue /= double(nr_gt);

    if(returnErrorMap) {
        int leakage;
        pListError = PyList_New(nr_pixel);
        for(int i = 0; i < nr_pixel; ++i) {
            leakage = accu_sp_size[GT[i]] - gt_size[GT[i]];
            pItem = PyFloat_FromDouble(double(leakage));
            PyList_SetItem(pListError, i, pItem);
        }
        pOutputObject = Py_BuildValue("dO", ue, pListError);
    }
    else {
        pOutputObject = Py_BuildValue("d", ue);
    }

    delete[] SP;
    delete[] GT;
    delete[] gt_hist;
    delete[] gt_size;
    delete[] sp_size;
    delete[] accu_sp_size;

    return pOutputObject;
}

static PyObject* computeSegCatCounts(PyObject* self, PyObject* args) {
    PyObject* pListSP;      // input
    PyObject* pListGT;      // input
    PyObject* pListOverlap;   // output
    PyObject* pOutputObject;   // output
    PyObject* pItem;
    double thresh = 0.2;
    int returnOverlapMap = 0;
    Py_ssize_t nr_pixel_sp;
    Py_ssize_t nr_pixel_gt;

    // Parse Python args
    if(!PyArg_ParseTuple(args, "OO|di", &pListSP, &pListGT, &thresh, &returnOverlapMap)) {
        PyErr_SetString(PyExc_TypeError, "Arg: SP_label_list, GT_label_list, thresh=0.2, returnOverlapMap=0");
        return NULL;
    }

    // Get input image from Python list to a double array
    nr_pixel_sp = PyList_Size(pListSP);
    nr_pixel_gt = PyList_Size(pListGT);
    if(nr_pixel_sp != nr_pixel_gt) {
        PyErr_SetString(PyExc_TypeError, "Number of pixels mismatched");
        printf("nr_pixel_sp = %d\n", int(nr_pixel_sp));
        printf("nr_pixel_gt = %d\n", int(nr_pixel_gt));
        return NULL;
    }

    int nr_pixel = int(nr_pixel_sp);

    int* SP = new int[nr_pixel];
    int* GT = new int[nr_pixel];

    for(int i = 0; i < nr_pixel; ++i) {
        pItem = PyList_GetItem(pListSP, i);
        SP[i] = PyInt_AsLong(pItem);
        pItem = PyList_GetItem(pListGT, i);
        GT[i] = PyInt_AsLong(pItem);
    }

    // Find number of superpixels and groundtruth segments: nr_sp, nr_gt
    int nr_sp = 0;
    int nr_gt = 0;
    for(int i = 0; i < nr_pixel; ++i) {
        if(SP[i] > nr_sp) nr_sp = SP[i];
        if(GT[i] > nr_gt) nr_gt = GT[i];
    }
    // include the starting point 0, so ++nr_sp
    ++nr_sp;
    ++nr_gt;

    int nr_one_to_one = 0;
    int nr_underseg = 0;
    int nr_overseg = 0;
    int* gt_hist = new int[nr_gt*nr_sp];    // a 2D table used to record SP and GT overlaps
    int* gt_size = new int[nr_gt];
    int* sp_size = new int[nr_sp];
    int* max_overlap_size = new int[nr_gt];
    double* max_overlap_perc = new double[nr_gt];
    int* min_sp_size_for_overlap = new int[nr_gt];
    double* sp_overlap_perc_on_gt = new double[nr_gt];
    memset(gt_hist, 0, sizeof(int)*nr_gt*nr_sp);
    memset(gt_size, 0, sizeof(int)*nr_gt);
    memset(sp_size, 0, sizeof(int)*nr_sp);
    for(int i = 0; i < nr_pixel; ++i) {
        gt_hist[GT[i]*nr_sp + SP[i]] += 1;
        gt_size[GT[i]] += 1;
        sp_size[SP[i]] += 1;
    }
    for(int i = 0; i < nr_gt; ++i) {
        max_overlap_size[i] = 0;
        min_sp_size_for_overlap[i] = -1;
        for(int j = 0; j < nr_sp; ++j) {
            if(gt_hist[i*nr_sp+j] > max_overlap_size[i]) {
                max_overlap_size[i] = gt_hist[i*nr_sp+j];
                min_sp_size_for_overlap[i] = sp_size[j];
            }
            else if (gt_hist[i*nr_sp+j] == max_overlap_size[i]) {
                max_overlap_size[i] = gt_hist[i*nr_sp+j];
                if (min_sp_size_for_overlap[i] == -1) {
                    min_sp_size_for_overlap[i] = sp_size[j];
                }
                else if (min_sp_size_for_overlap[i] > sp_size[j]) {
                    min_sp_size_for_overlap[i] = sp_size[j];
                }
            }
        }
        sp_overlap_perc_on_gt[i] = double(min_sp_size_for_overlap[i]) / double(gt_size[i]);
        max_overlap_perc[i] = double(max_overlap_size[i]) / double(gt_size[i]);
//        cout << "gt_size[i] : " << gt_size[i] << endl;
//        cout << "max_overlap_size[i] : " << max_overlap_size[i] << endl;
//        cout << "min_sp_size_for_overlap[i] : " << min_sp_size_for_overlap[i] << endl;
//        cout << "max_overlap_perc[i]: " << max_overlap_perc[i] << endl;
//        cout << "sp_overlap_perc_on_gt[i] : " << sp_overlap_perc_on_gt[i] << endl;

        if(max_overlap_perc[i] < 1.0 - thresh) {
            nr_overseg += 1;
        }
        else if((max_overlap_perc[i] >= 1.0 - thresh) && (sp_overlap_perc_on_gt[i] <= 1.0 + thresh)) {
            nr_one_to_one += 1;
        }
        else {
            nr_underseg += 1;
        }
    }

    if(returnOverlapMap) {
        pListOverlap = PyList_New(nr_pixel);
        for(int i = 0; i < nr_pixel; ++i) {
            pItem = PyFloat_FromDouble(max_overlap_perc[GT[i]]);
            PyList_SetItem(pListOverlap, i, pItem);
        }
        pOutputObject = Py_BuildValue("iiiiO", nr_underseg, nr_one_to_one, nr_overseg, nr_gt, pListOverlap);
    }
    else {
        pOutputObject = Py_BuildValue("iiii", nr_underseg, nr_one_to_one, nr_overseg, nr_gt);
    }

    delete[] SP;
    delete[] GT;
    delete[] gt_hist;
    delete[] gt_size;
    delete[] sp_size;
    delete[] max_overlap_size;
    delete[] max_overlap_perc;
    delete[] min_sp_size_for_overlap;
    delete[] sp_overlap_perc_on_gt;

    return pOutputObject;
}

static PyObject* computeBR(PyObject* self, PyObject* args) {
    PyObject* pListSP;      // input
    PyObject* pListGT;      // input
    PyObject* pOutputObject;   // output
    PyObject* pItem;
    Py_ssize_t nr_pixel_sp;
    Py_ssize_t nr_pixel_gt;

    int h, w, nr_pixel;
    int r = 1;

    // Parse Python args
    if(!PyArg_ParseTuple(args, "OOii|i", &pListSP, &pListGT, &h, &w, &r)) {
        PyErr_SetString(PyExc_TypeError, "Arg: SP_label_list, GT_label_list, height, width, r=1");
        return NULL;
    }

    // Get input image from Python list to a double array
    nr_pixel = h*w;
    nr_pixel_sp = PyList_Size(pListSP);
    nr_pixel_gt = PyList_Size(pListGT);
    if(nr_pixel_sp != nr_pixel_gt) {
        PyErr_SetString(PyExc_TypeError, "Number of pixels mismatched");
        printf("nr_pixel_sp = %d\n", int(nr_pixel_sp));
        printf("nr_pixel_gt = %d\n", int(nr_pixel_gt));
        return NULL;
    }
    if(nr_pixel != nr_pixel_gt) {
        PyErr_SetString(PyExc_TypeError, "Size mismatched");
        printf("nr_pixel_sp = %d\n", int(nr_pixel_sp));
        printf("h*w = %d\n", nr_pixel);
        return NULL;
    }

    int* SP = new int[nr_pixel];
    int* GT = new int[nr_pixel];

    for(int i = 0; i < nr_pixel; ++i) {
        pItem = PyList_GetItem(pListSP, i);
        SP[i] = int(PyInt_AsLong(pItem));
        pItem = PyList_GetItem(pListGT, i);
        GT[i] = int(PyInt_AsLong(pItem));
    }

    unsigned char* transitionSP = new unsigned char[nr_pixel];
    unsigned char* transitionGT = new unsigned char[nr_pixel];
    memset(transitionSP, 0, sizeof(unsigned char)*nr_pixel);
    memset(transitionGT, 0, sizeof(unsigned char)*nr_pixel);
    int index = 0;
    for(int y = 0; y < h; ++y) {
        for(int x = 0; x < w; ++x) {
            if(x < w-1) {
                if(SP[index] != SP[index+1]) transitionSP[index] = 1;
                if(GT[index] != GT[index+1]) transitionGT[index] = 1;
            }
            if(y < h-1) {
                if(SP[index] != SP[index+w]) transitionSP[index] = 1;
                if(GT[index] != GT[index+w]) transitionGT[index] = 1;
            }
            ++index;
        }
    }

    // Calculate boundary recall on the ground truth boundary/transition map
    int nr_transition_gt = 0;
    int nr_true_positive_gt = 0;
    bool hit_gt;
    int index_gt = 0;
    double minimal_error_pixel_gt;
    double tmp_error_pixel_gt;
    double total_error_pixel_gt = 0.0;
    for(int y = 0; y < h; ++y) {
        for(int x = 0; x < w; ++x) {
            if(transitionGT[index_gt]) {
                ++nr_transition_gt;
                hit_gt = false;
                minimal_error_pixel_gt = -1.0;
                for(int u = -r; u <= r; ++u) {
                    for(int v = -r; v <= r; ++v) {
                        if((y+u >= 0) && (y+u < h) && (x+v >= 0) && (x+v < w)) {
                            if(transitionSP[index_gt+u*w+v] > 0) {
                                hit_gt = true;
                                tmp_error_pixel_gt = sqrt(double(pow(u, 2) + pow(v, 2)));
                                if(minimal_error_pixel_gt == -1.0) {
                                    minimal_error_pixel_gt = tmp_error_pixel_gt;
                                }
                                else if(minimal_error_pixel_gt > tmp_error_pixel_gt) {
                                    minimal_error_pixel_gt = tmp_error_pixel_gt;
                                }
//                                break;
                            }
                        }
                    }
//                    if(hit_gt) break;
                }
//                cout << "minimal_error_pixel_gt (outside) : " << minimal_error_pixel_gt << endl;
                if(hit_gt) ++nr_true_positive_gt;
                if(minimal_error_pixel_gt >= 0.0) {
                    total_error_pixel_gt += minimal_error_pixel_gt;
                }
            }
            ++index_gt;
        }
    }

    // Calculate boundary recall on the superpixel boundary/transition map
    int nr_transition_sp = 0;
    int nr_true_positive_sp = 0;
    bool hit_sp;
    int index_sp = 0;
    double minimal_error_pixel_sp;
    double tmp_error_pixel_sp;
    double total_error_pixel_sp = 0.0;
    for(int y = 0; y < h; ++y) {
        for(int x = 0; x < w; ++x) {
            if(transitionSP[index_sp]) {
                ++nr_transition_sp;
                hit_sp = false;
                minimal_error_pixel_sp = -1.0;
                for(int u = -r; u <= r; ++u) {
                    for(int v = -r; v <= r; ++v) {
                        if((y+u >= 0) && (y+u < h) && (x+v >= 0) && (x+v < w)) {
                            if(transitionGT[index_sp+u*w+v] > 0) {
                                hit_sp = true;
                                tmp_error_pixel_sp = sqrt(double(pow(u, 2) + pow(v, 2)));
                                if(minimal_error_pixel_sp == -1.0) {
                                    minimal_error_pixel_sp = tmp_error_pixel_sp;
                                }
                                else if(minimal_error_pixel_sp > tmp_error_pixel_sp) {
                                    minimal_error_pixel_sp = tmp_error_pixel_sp;
                                }
//                                break;
                            }
                        }
                    }
//                    if(hit_sp) break;
                }
//                cout << "minimal_error_pixel_sp (outside) : " << minimal_error_pixel_sp << endl;
                if(hit_sp) ++nr_true_positive_sp;
                if(minimal_error_pixel_sp >= 0.0) {
                    total_error_pixel_sp += minimal_error_pixel_sp;
                }
            }
            ++index_sp;
        }
    }

    double br_gt;
    double br_sp;
    double average_error_pixel_gt;
    double average_error_pixel_sp;
    if(nr_transition_gt && nr_transition_sp) {
        br_gt = double(nr_true_positive_gt)/double(nr_transition_gt);
        br_sp = double(nr_true_positive_sp)/double(nr_transition_sp);
        if (nr_true_positive_gt) {
            average_error_pixel_gt = total_error_pixel_gt / double(nr_true_positive_gt);
        }
        else {
            average_error_pixel_gt = -1.0;
        }
        if (nr_true_positive_sp) {
            average_error_pixel_sp = total_error_pixel_sp / double(nr_true_positive_sp);
        }
        else {
            average_error_pixel_sp = -1.0;
        }
    }
    else if(nr_transition_gt) {
        br_gt = 0.0;
        average_error_pixel_gt = -1.0;
        br_sp = 1.0;
        average_error_pixel_sp = 0.0;
    }
    else if(nr_transition_sp) {
        br_gt = 1.0;
        average_error_pixel_gt = 0.0;
        br_sp = 0.0;
        average_error_pixel_sp = -1.0;
    }
    else {
        br_gt = 1.0;
        average_error_pixel_gt = 0.0;
        br_sp = 1.0;
        average_error_pixel_sp = 0.0;
    }

    pOutputObject = Py_BuildValue("dddd", br_gt, br_sp, average_error_pixel_gt, average_error_pixel_sp);

    delete[] SP;
    delete[] GT;
    delete[] transitionSP;
    delete[] transitionGT;

    return pOutputObject;
}

static PyMethodDef EvalSPMethods[] = {
    {
        "computeASA",
        computeASA,
        METH_VARARGS,
        ""
    },
    {
        "computeUE",
        computeUE,
        METH_VARARGS,
        ""
    },
    {
        "computeSegCatCounts",
        computeSegCatCounts,
        METH_VARARGS,
        ""
    },
    {
        "computeBR",
        computeBR,
        METH_VARARGS,
        ""
    },
    {NULL, NULL, 0, NULL}
};

#if PY_MAJOR_VERSION >= 3
static struct PyModuleDef EvalSPModule = {
    PyModuleDef_HEAD_INIT,
    "EvalSPModule",
    NULL,
    -1,
    EvalSPMethods
};
#endif

#if PY_MAJOR_VERSION >= 3
PyMODINIT_FUNC PyInit_EvalSPModule(void) {
     return PyModule_Create(&EvalSPModule);
}
#else
PyMODINIT_FUNC initEvalSPModule(void) {
    (void)Py_InitModule("EvalSPModule", EvalSPMethods);
}
#endif
